batch_size = 1
input_size = 4
index_chart = ['e', 'h', 'l', 'o']  
# 预测的结果是设置的索引,这样方便对结果更快地进行取值
x_data = torch.LongTensor([[1, 0, 2, 2, 3]]).view(5, 1)
print(x_data.shape)
y_data = [3, 1, 2, 3, 2]  # 标签
labels = torch.LongTensor(y_data).view(-1, 1)  # 调整维度
emb = nn.Embedding(4, 10)
inputs = emb(x_data) # 将4维输入数据嵌入成10维,为后面的linear_ix等操作进行设置(5,1,10),嵌入值均是随机设定
print(inputs)
print(inputs.shape)

class emb_lstm(nn.Module):
    def __init__(self):
        super(emb_lstm, self).__init__()
        # 上一时刻的内部状态g(x),外部状态h(x)
        
        # 输入门,输入值和上一阶段的外部状态值通过激活函数后成为输入值it
        self.linear_ix = nn.Linear(10, 4) 
        # 遗忘门,输入值和上一阶段的外部状态值通过激活函数后成为遗忘值ft
        self.linear_fx = nn.Linear(10, 4) 
        # 候选变量,输入值和上一阶段的外部状态值通过Tanh激活函数后成为候选变量值
        self.linear_gx = nn.Linear(10, 4) 
         #输出值,输入值和上一阶段的外部状态值通过激活函数激活后成为输出值ot
        self.linear_ox = nn.Linear(10, 4)
        # 隐藏层的LSTM变换
        self.linear_ih = nn.Linear(4, 4)
        self.linear_fh = nn.Linear(4, 4)
        self.linear_gh = nn.Linear(4, 4)
        self.linear_oh = nn.Linear(4, 4)
        self.Sigmoid = nn.Sigmoid()
        self.Tanh = nn.Tanh()

    def forward(self, x, hidden, c): 
        # 输入x,外部状态hidden,以及候选状态c
        # 输入值x和外部状态h(x)相结合,
        # 再通过激活函数激活得到内部状态的i,f,g(候选状态),o值；
        i = self.Sigmoid(self.linear_ix(x) + self.linear_ih(hidden)) 
        # i:(1,4)
        f = self.Sigmoid(self.linear_fx(x) + self.linear_fh(hidden)) 
        # f:(1,4)
        g = self.Tanh(self.linear_gx(x) +  self.linear_gh(hidden)) 
        # g:(1,4)
        o = self.Sigmoid(self.linear_ox(x) + self.linear_oh(hidden)) 
        # o:(1,4) 

        # 候选状态g乘以输入值i,
        # 再加上一时刻的内部状态c乘以遗忘值f,
        # 得到该时刻的更新的内部状态值c     
        c = f * c + i * g   

        # 输出元素o乘以经过激活后的该时刻的内部状态值,
        # 得到该时刻的外部状态值
        hidden = o * self.Tanh(c)  # 上式得到的结果c通过输出门
        return hidden, c

model = emb_lstm()
# 构建损失函数和优化器方法
criterion = nn.CrossEntropyLoss()  # 交叉熵
optimizer = torch.optim.Adam(model.parameters(), lr=0.06)
# 进行训练和更新
for epoch in range(100):
    loss = 0
    optimizer.zero_grad()
    hidden = torch.zeros(batch_size, input_size)  
    # 提供初始化隐藏层

    c = torch.zeros(batch_size, input_size)
    print('Predicten string:', end='') 

    #并行遍历数据集 一个一个训练
    for input, label in zip(inputs, labels): 
        hidden, c = model(input, hidden, c) # input:(1,10)
        loss += criterion(hidden, label)
        _, idx = hidden.max(dim=1) 
        #从第一个维度(即行维度)上取出预测概率最大的值和该值所在序号
        
        print(index_chart[idx.item()], end='') 
        #按上面序号输出相应字母字符

    loss.backward(retain_graph=True)
    optimizer.step()
    print(', Epoch [%d/100] loss=%.4f' % (epoch + 1, loss.item())) 
